# LeetCode 23. 合并K个升序链表
# https://leetcode.cn/problems/merge-k-sorted-lists/description/

from typing import List, Optional
import heapq
# Definition for singly-linked list.
class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next


class Solution:
    def mergeKLists(self, lists: List) -> Optional[ListNode]:
        result = ListNode()
        curr_result = []
        for sub_list in lists:
            while sub_list:
                heapq.heappush(curr_result, sub_list.val)
                sub_list = sub_list.next

        head = result
        while curr_result:
            num = heapq.heappop(curr_result)
            curr_node = ListNode(num, None)
            result.next = curr_node
            result = result.next
        return head.next


class Solution2:
    def mergeKLists(self, lists: List) -> Optional[ListNode]:
        ListNode.__lt__ = lambda a, b: 0
        if not lists:
            return None
        result = ListNode()
        head = result
        heap = []
        for node in lists:
            if node:
                heapq.heappush(heap, (node.val, node))

        while heap:
            num, node = heapq.heappop(heap)
            result.next = ListNode(num, None)
            result = result.next
            if node.next:
                node = node.next
                heapq.heappush(heap, (node.val, node))
        return head.next


if __name__ == '__main__':
    lists = [[1, 4, 5], [1, 3, 4], [2, 6]]
    # result = Solution().mergeKLists(lists)
    # print(result)
    result2 = Solution2().mergeKLists(lists)
    print(result2)
